from torchvision import transforms
from torch.utils.data import DataLoader, Dataset
from PIL import Image
import glob

label_name = ["airplane",
              "automobile",
              "bird",
              "cat",
              "deer",
              "dog",
              "frog",
              "horse",
              "ship",
              "truck"]
label_dict = {}

for idx, name in enumerate(label_name):
    label_dict[name] = idx


def default_loader(path):
    return Image.open(path).convert("RGB")


train_transform = transforms.Compose([
    transforms.RandomResizedCrop((28, 28)),
    transforms.RandomHorizontalFlip(),
    transforms.RandomVerticalFlip(),
    # transforms.RandomRotation(90),
    # transforms.RandomGrayscale(0.1),
    # transforms.ColorJitter(0.3, 0.3, 0.3, 0.3),
    # transforms.RandomCrop(28),
    # transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))
])

test_transform = transforms.Compose([
    transforms.CenterCrop((28, 28)),
    transforms.ToTensor(),
    transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))
])


# 定义类 继承Dataset
class MyDataset(Dataset):
    # 数据读取
    def __init__(self, im_list, transform=None, loader=default_loader):
        # 初始化类
        super(MyDataset, self).__init__()
        images = []
        for im_item in im_list:
            im_label_name = im_item.split("\\")[-2]
            images.append([im_item, label_dict[im_label_name]])

        self.images = images
        self.transform = transform
        self.loader = loader

    # getitem 中对路径数据加载 并增强转换为张量
    def __getitem__(self, index):
        im_path, im_label = self.images[index]
        im_data = self.loader(im_path)
        if self.transform is not None:
            im_data = self.transform(im_data)
        return im_data, im_label

    # 大小
    def __len__(self):
        return len(self.images)


im_train_list = glob.glob("E:\pythonProject\dataset\cifar10\cifar-10-batches-py\TRAIN\*\*.png")
im_test_list = glob.glob("E:\pythonProject\dataset\cifar10\cifar-10-batches-py\TEST\*\*.png")
train_dataset = MyDataset(im_train_list, transform=train_transform)
test_dataset = MyDataset(im_test_list, transform=test_transform)
train_data_loader = DataLoader(dataset=train_dataset, batch_size=128, shuffle=True, num_workers=4)
test_data_loader = DataLoader(dataset=test_dataset, batch_size=128, shuffle=False, num_workers=4)
#
# print("mun_of_train", len(train_dataset))
# print("mun_of_test", len(test_dataset))
